{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verilator 代码生成测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from infrastructure import (\n",
    "    compile_hardware,\n",
    "    compiler_opts,\n",
    "    run_module,\n",
    "    offload,\n",
    "    clear_stats,\n",
    "    stats\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relay\n",
    "\n",
    "def create_module_add(shape, dtype):\n",
    "    \"\"\"Create add module.\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    shape : Tuple\n",
    "        The shape tuple.\n",
    "\n",
    "    dtype : Str\n",
    "        The data type.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mod: Module\n",
    "        The relay module.\n",
    "    \"\"\"\n",
    "    x = relay.var(\"x\", shape=shape, dtype=dtype)\n",
    "    y = relay.var(\"y\", shape=shape, dtype=dtype)\n",
    "    z = relay.add(x, y)\n",
    "    f = relay.Function([x, y], z)\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = f\n",
    "    return mod\n",
    "\n",
    "\n",
    "def create_module_bias_add(xshape, yshape, dtype):\n",
    "    \"\"\"Create bias_add module.\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    xshape : Tuple\n",
    "        The x shape tuple.\n",
    "\n",
    "    yshape : Tuple\n",
    "        The y shape tuple.\n",
    "\n",
    "    dtype : Str\n",
    "        The data type.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mod: Module\n",
    "        The relay module.\n",
    "    \"\"\"\n",
    "    x = relay.var(\"x\", shape=xshape, dtype=dtype)\n",
    "    y = relay.var(\"y\", shape=yshape, dtype=dtype)\n",
    "    z = relay.nn.bias_add(x, y, axis=3)\n",
    "    f = relay.Function([x, y], z)\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = f\n",
    "    return mod\n",
    "\n",
    "\n",
    "def run_and_check(xshape, yshape, dtype, mod, opts):\n",
    "    \"\"\"Run and check values.\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    xshape : Tuple\n",
    "        The x shape tuple.\n",
    "\n",
    "    yshape : Tuple\n",
    "        The y shape tuple.\n",
    "\n",
    "    dtype : Str\n",
    "        The data type.\n",
    "\n",
    "    mod: Module\n",
    "        The relay module.\n",
    "\n",
    "    opts: Dict\n",
    "        The compiler options.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    cycles: Int\n",
    "        The number of cycles.\n",
    "    \"\"\"\n",
    "    x_data = np.random.randint(5, size=xshape, dtype=dtype)\n",
    "    y_data = np.random.randint(5, size=yshape, dtype=dtype)\n",
    "    ref = x_data + y_data\n",
    "    inp = {\"x\": x_data, \"y\": y_data}\n",
    "    clear_stats()\n",
    "    out = run_module(inp, mod, params=None, opts=opts)\n",
    "    values = stats()\n",
    "    tvm.testing.assert_allclose(out.numpy(), ref, rtol=1e-5, atol=1e-5)\n",
    "    return values[\"cycle_counter\"]\n",
    "\n",
    "\n",
    "def print_test_info(test, lanes, cycles):\n",
    "    \"\"\"Print counter\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    test : Str\n",
    "        The name of the test.\n",
    "\n",
    "    lanes : Int\n",
    "        The number of vector lanes.\n",
    "\n",
    "    cycles : Int\n",
    "        The number of cycles.\n",
    "    \"\"\"\n",
    "    print(\"test:{} vector-lanes:{} number of cycles:{}\".format(test, lanes, cycles))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tadd(lanes, verilator_app_path):\n",
    "    \"\"\"Print counter\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    lanes : Int\n",
    "        The number of vector lanes.\n",
    "    \"\"\"\n",
    "    dtype = \"int32\"\n",
    "    shape = (8, 4)\n",
    "    mod = create_module_add(shape, dtype)\n",
    "    mod = offload(mod)\n",
    "    lib = compile_hardware(lanes, verilator_app_path)\n",
    "    opts = compiler_opts(lib)\n",
    "    cycles = run_and_check(shape, shape, dtype, mod, opts)\n",
    "    print_test_info(\"add\", lanes, cycles)\n",
    "\n",
    "def tbias(lanes, verilator_app_path):\n",
    "    \"\"\"Print counter\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    lanes : Int\n",
    "        The number of vector lanes.\n",
    "    \"\"\"\n",
    "    dtype = \"int32\"\n",
    "    xshape = (1, 112, 112, 32)\n",
    "    yshape = (32,)\n",
    "    mod = create_module_bias_add(xshape, yshape, dtype)\n",
    "    mod = offload(mod)\n",
    "    lib = compile_hardware(lanes, verilator_app_path)\n",
    "    opts = compiler_opts(lib)\n",
    "    cycles = run_and_check(xshape, yshape, dtype, mod, opts)\n",
    "    print_test_info(\"nn.bias_add\", lanes, cycles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tbias(lanes, verilator_app_path):\n",
    "    \"\"\"Print counter\n",
    "\n",
    "    Paramters\n",
    "    ---------\n",
    "    lanes : Int\n",
    "        The number of vector lanes.\n",
    "    \"\"\"\n",
    "    dtype = \"int32\"\n",
    "    xshape = (1, 112, 112, 32)\n",
    "    yshape = (32,)\n",
    "    mod = create_module_bias_add(xshape, yshape, dtype)\n",
    "    mod = offload(mod)\n",
    "    lib = compile_hardware(lanes, verilator_app_path)\n",
    "    opts = compiler_opts(lib)\n",
    "    cycles = run_and_check(xshape, yshape, dtype, mod, opts)\n",
    "    print_test_info(\"nn.bias_add\", lanes, cycles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[15:23:02] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test:add vector-lanes:1 number of cycles:32\n",
      "test:add vector-lanes:4 number of cycles:8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[15:23:05] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n"
     ]
    }
   ],
   "source": [
    "from set_env import TVM_ROOT\n",
    "# add tests\n",
    "verilator_app_path = f\"{TVM_ROOT}/3rdparty/vta-hw/apps/verilator/add\"\n",
    "tadd(1, verilator_app_path)\n",
    "tadd(4, verilator_app_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[15:23:05] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test:nn.bias_add vector-lanes:1 number of cycles:401408\n",
      "test:nn.bias_add vector-lanes:32 number of cycles:12544\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[15:23:08] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module\n"
     ]
    }
   ],
   "source": [
    "# bias_add tests\n",
    "tbias(1, verilator_app_path)\n",
    "tbias(32, verilator_app_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
